#include <chroma.h>
#include <actions/ferm/invert/syssolver_linop_cg_array.h>
#include "dwf_4d_invert.h"

enum { DaggerNo, DaggerYes };
using namespace Chroma;

typedef LatticeFermion T;
typedef multi1d<LatticeFermion> T5;
typedef multi1d<LatticeColorMatrix> U;

class qudaDwf4d {
    public:
        int Ls;
        double mass;
        double M5;
        double residual;
        int max_iter;
        multi1d<LatticeColorMatrix> gauge;

        Handle< FermBC<T,U,U> > *fbc;
        Handle<CreateFermState<T,U,U> > *cfs;
        UnprecDWFermActArray  *S_f;
        Handle< FermState<T,U,U> > *fs;
        Handle< UnprecLinearOperatorArray<T,U,U> > *M5ddwf;




        void end(void);
        void init(double m, double M5, int N5, multi1d<LatticeColorMatrix> U, int dev);
        double Munprec(const LatticeFermion &chi,
                LatticeFermion &psi, 
                int dag);
        int CGNE_MdagM(LatticeFermion &sol,const LatticeFermion &src);
        int CGNE_Mdag (LatticeFermion &sol,const LatticeFermion &src);
        int CGNE_M    (LatticeFermion &sol,const LatticeFermion &src);

};
void qudaDwf4d::end(void)
{
    endQuda();
    delete fbc, cfs, S_f, fs, M5ddwf;
}
void qudaDwf4d::init(double m, double m5, int N5, multi1d<LatticeColorMatrix> gaugeU, int device)
{
    mass = m;
    Ls = N5;
    gauge = gaugeU;
    M5 = m5;

    residual = 1.0e-8;
    max_iter = 1000;



    // Initialise QUDA and QUDADWF here
    initQuda(device);

    multi1d<int> bcs(Nd);
    bcs[0] = bcs[1] = bcs[2] = bcs[3] = 1;


    fbc = new Handle< FermBC<T,U,U> > (new SimpleFermBC< T, U, U >(bcs));

    cfs = new Handle<CreateFermState<T,U,U> > ( new CreateSimpleFermState<T,U,U>(*fbc));
    S_f = new UnprecDWFermActArray  (*cfs, M5, mass, Ls);
    fs = new Handle< FermState<T,U,U> > ( (*S_f).createState(gaugeU) );

    M5ddwf = new Handle< UnprecLinearOperatorArray<T,U,U> > ((*S_f).unprecLinOp(*fs,mass));




    /*
    //Physics parameters
    dwf.solver = DWF;
    dwfa.Ls           = Ls;
    dwfa.M5           = M5;
    dwfa.mass         = mass;
    dwfa.precon_5d    = 1;
    dwfa.max_iter     = 1000;
    dwfa.residual     = 1.0e-8;
    //Geometry
    dwfa.node_latt[0] = lx;
    dwfa.node_latt[1] = ly;
    dwfa.node_latt[2] = lz;
    dwfa.node_latt[3] = lt;

    for(int mu=0;mu<4;mu++){
    dwfa.local_comm[mu] = 1;
    }

    dop.init(dwfa);
    dop.importGauge(U);
     */
}

int qudaDwf4d::CGNE_M(LatticeFermion &psi,const LatticeFermion &src)
{
    LatticeFermion MdagSrc;
    Munprec(src,MdagSrc,DaggerYes);
    return CGNE_MdagM(psi,MdagSrc);
}
int qudaDwf4d::CGNE_Mdag(LatticeFermion &psi,const LatticeFermion &src)
{
    LatticeFermion MinvPsi;
    int iter = CGNE_MdagM(MinvPsi,src);
    Munprec(MinvPsi,psi,DaggerNo);
}
int qudaDwf4d::CGNE_MdagM(LatticeFermion &psi,const LatticeFermion &src)
{
    LatticeFermion p;
    LatticeFermion mp;
    LatticeFermion mmp;
    LatticeFermion r;

    Munprec(psi,mp,DaggerNo);
    Munprec(mp,mmp,DaggerYes);

    r = src - mmp;
    p = r;

    Real ssq = norm2(src);
    Real rsq =  residual* residual*ssq;
    Real cp  = norm2(r);

    if ( toDouble(cp) <= toDouble(rsq) ) {
        QDPIO::cout <<"CGNE_unprec converged k=0 - nice guess "<<cp << " " << rsq<<endl;
        return 0;
    }  

    Real c,a,d,b;

    for (int k=1;k<=max_iter;k++){

        c=cp;

        d = Munprec(p,mp,DaggerNo);

        // Update solution vector
        a = c/d;
        psi = a*p+psi;

        // MdagMp
        Munprec(mp,mmp,DaggerYes); 

        r = r-a*mmp;
        cp= norm2(r);
        if ( 0 ) 
            QDPIO::cout <<"CGNE_unprec k="<<k<<" r^2 = "<<cp<<endl;

        // Stopping condition
        if ( toDouble(cp) <= toDouble(rsq) ) { 
            QDPIO::cout <<"CGNE_unprec converged k="<<k<<" "<<cp << " " << rsq<<endl;
            Munprec(psi,mp,DaggerNo);
            Munprec(mp,mmp,DaggerYes); 
            r=mmp-src;
            Real true_residual = sqrt(norm2(r)/norm2(src));
            QDPIO::cout <<"CGNE_unprec true residual "<<true_residual<<endl;
            return k;
        }

        // New (conjugate/M-orthogonal) search direction
        b = cp/c;
        p = b*p+r;
    }
    QDPIO::cout <<"CGNE_unprec not converged"<<endl;
    return -1;
}
double qudaDwf4d::Munprec(const LatticeFermion &chi,
        LatticeFermion &psi, 
        int dag)
{
    multi1d<LatticeFermion> chi_5(Ls), psi_5(Ls), tmp(Ls);

    /**************************************************************
     * 4d -> 5d system
     **************************************************************/
    // Need chi_5[j] = P_[j][0] chi
    // Seems screwey compared to physical field conventions in DWF.
    // Perhaps wrong?
    chi_5[0]   = chiralProjectMinus(chi);
    chi_5[Ls-1]= chiralProjectPlus(chi);

    // Zero bulk of chi and result psi
    for(int s=1;s<Ls-1;s++) chi_5[s]=zero;
    for(int s=0;s<Ls;s++)   psi_5[s]=zero;

    if ( dag ) { 
        /**************************************************************
         * PV^{-d}
         **************************************************************/
        dwf_4d_invert::fullInversion( chi_5, tmp, gauge, M5, DaggerYes);

        /**************************************************************
         * D(m)
         **************************************************************/

        (**M5ddwf)(psi_5, tmp, MINUS); // Chroma

    } else { 
        /**************************************************************
         * D(m)
         **************************************************************/

        (**M5ddwf)(tmp, chi_5, PLUS); // Chroma

        /**************************************************************
         * PV^{-1}
         **************************************************************/
        dwf_4d_invert::fullInversion( tmp, psi_5, gauge, M5, DaggerNo);
    }

    /********************************************************
     * Merge the chiralities from the walls to reconstruct 4d prop
     ********************************************************/

    // Need psi = P^{inv}[0][j] psi_5[j]
    psi = chiralProjectMinus(psi_5[0]);
    psi += chiralProjectPlus (psi_5[Ls-1]);

    double nrm = toDouble(norm2(psi));
    return nrm;
}

#define Printf printf

int main (int argc,char **argv )
{
    Chroma::initialize(&argc,&argv);

    /********************************************************
     * Setup QDP
     ********************************************************
     */
    multi1d<int> nrow(Nd);
    nrow[0] = 4;
    nrow[1] = 4;
    nrow[2] = 4;
    nrow[3] = 4;
    int Ls = 8;
    double mass =0.1;
    double m5   =1.8;

    Layout::setLattSize(nrow);
    Layout::create();

    /********************************************************
     * Gaussian gauge field
     ********************************************************
     */

    multi1d<LatticeColorMatrix>  u(Nd); HotSt(u);

    /********************************************************
     * Gaussian source and result vectors
     ********************************************************
     */
    LatticeFermion psi;
    LatticeFermion chi;
    LatticeFermion psi_qdp;
    LatticeFermion chi_qdp;

    /* 
     * Chroma regression test set up
     */
    // BCs and gauge field
    multi1d<int> bcs(Nd);
    bcs[0] = bcs[1] = bcs[2] = bcs[3] = 1;
    Handle< FermBC<T,U,U> > fbc(new SimpleFermBC< T, U, U >(bcs));
    Handle<CreateFermState<T,U,U> > cfs( new CreateSimpleFermState<T,U,U>(fbc));

    // 5d DWF even-odd operator
    EvenOddPrecDWFermActArray  DWF_5d(cfs, m5, mass, Ls);
    Handle< FermState<T,U,U> >                     fs_5d( DWF_5d.createState(u) );
    Handle< EvenOddPrecLinearOperatorArray<T,U,U> >  M_5d(DWF_5d.precLinOp(fs_5d,mass));
    Handle< EvenOddPrecLinearOperatorArray<T,U,U> > PV_5d(DWF_5d.precLinOp(fs_5d,1.0));

    SysSolverCGParams invParam;
    invParam.RsdCG        = 1.0e-8;
    invParam.RsdCGRestart = 1.0e-8;
    invParam.MaxCG        = 1000;
    invParam.MaxCGRestart = 1000;

    // 4d effective overlap operator
    UnprecPPDWF4DLinOp<T,U,U>  Overlap_4d(M_5d,PV_5d,invParam);

    QDPIO::cout << "Initialising new linop" <<endl;
    {
        qudaDwf4d linop;
        linop.init(mass,m5,Ls,u,0);

        /*
         * Forward matrix
         */
        gaussian(psi);
        gaussian(psi_qdp);
        gaussian(chi);

        QDPIO::cout << "Applying operator" <<endl;
        linop.Munprec(chi,psi,DaggerNo);

        QDPIO::cout << "Applying Chroma operator"<<endl;
        Overlap_4d(psi_qdp,chi,PLUS);
        QDPIO::cout << "|5dQuda-4dOp| = "<<norm2(psi-psi_qdp)<<endl;
        QDPIO::cout << "|4dOp| = "<<norm2(psi_qdp)<<endl;
        QDPIO::cout << "|5dQuda| = "<<norm2(psi)<<endl;

        gaussian(psi);
        gaussian(chi);
        QDPIO::cout << "Applying adjoint operator" <<endl;
        linop.Munprec(chi,psi,DaggerYes);

        QDPIO::cout << "Applying Chroma operator"<<endl;
        Overlap_4d(psi_qdp,chi,MINUS);
        QDPIO::cout << "|5dQuda-4dOp| = "<<norm2(psi-psi_qdp)<<endl;
        QDPIO::cout << "|4dOp| = "<<norm2(psi_qdp)<<endl;
        QDPIO::cout << "|5dQuda| = "<<norm2(psi)<<endl;


        
        QDPIO::cout << "Running CGNE for effective 4d operator" <<endl;
        psi=zero;
        gaussian(chi);
        linop.CGNE_M(psi,chi);

        //QDPIO::cout << "Subtracting contact term and normalising"<<endl;
        //psi = (psi - chi) *Real(1.0/(1.0-mass));
        
        QDPIO::cout << "Applying Chroma 4d Operator"<<endl;
        Overlap_4d(chi_qdp,psi,PLUS);

        QDPIO::cout << "|Source - DWF4d*PSI| = "<<norm2(chi-chi_qdp)<<endl;

        linop.end();
    }

    // Should be able to relate 4d effective operator to inverse + contact term
    // Think about this as a test.

}








